import torch
from torch.optim import AdamW

from MyData import MyDataset
from torch.utils.data import DataLoader
from Net import Model
from transformers import BertTokenizer

# 定义训练设备
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
EPOCHS = 100
BATCH_SIZE = 32
# 加载分词器
model_name = "./model/google-bert/bert-base-chinese/models--google-bert--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
token = BertTokenizer.from_pretrained(model_name)


# 自定义数据编码处理函数
def collate_fn(data):
    sente = [i[0] for i in data]
    label = [i[1] for i in data]
    # 编码处理
    data = token.batch_encode_plus(
        batch_text_or_text_pairs=sente,
        truncation=True,
        padding='max_length',
        max_length=300,
        return_tensors='pt',
        return_length=True
    )
    input_ids = data['input_ids']
    attention_mask = data['attention_mask']
    token_type_ids = data['token_type_ids']
    labels = torch.LongTensor(label)

    return input_ids, attention_mask, token_type_ids, labels


# 创建数据集
train_dataset = MyDataset("train")
# 创建数据加载器
train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    drop_last=True,
    collate_fn=collate_fn  # 只传递函数名
)

if __name__ == '__main__':
    # start train
    print(DEVICE)
    model = Model().to(DEVICE)
    # 优化器
    optimizer = AdamW(model.parameters(), lr=5e-4)
    # loss function
    loss_func=torch.nn.CrossEntropyLoss()
    # train model
    model.train()
    for epoch in range(EPOCHS):
        for i,(input_ids, attention_mask, token_type_ids, labels) in enumerate(train_loader):
            # 将数据放到DEVICE上
            input_ids, attention_mask, token_type_ids, labels=input_ids.to(DEVICE), attention_mask.to(DEVICE), token_type_ids.to(DEVICE), labels.to(DEVICE)
            # 执行前向计算
            out=model(input_ids, attention_mask, token_type_ids)
            # print("Output shape:", out.shape)
            # 计算损失
            loss = loss_func(out, labels)
            # 优化模型 清空权重 反向传播 更新梯度
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 每5个查看一下
            if i%20==0:
                out=out.argmax(dim=1)
                #  准确率
                acc=(out==labels).sum()/len(labels)
                print(epoch,i,loss.item(),acc)
        # 保存模型参数
        torch.save(model.state_dict(),f"./params/{epoch}bert.pt")
        print(epoch)